import collections
class Solution(object):
    def findDiagonalOrder(self, nums):
        d = collections.defaultdict(list)
        n = len(nums)
        for i in range(n):
            for j in range(len(nums[i])):
                d[i + j].append((i, j))
        k_list = list(d.keys())
        k_list.sort()
        out = []
        for k in k_list:
            d[k].sort(reverse=True)
            for x, y in d[k]:
                out.append(nums[x][y])
        return out
